package com.lanhao.aot.llm.doubao;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import com.alibaba.fastjson.JSON;
import com.lanhao.aot.llm.dto.AIChatDTO;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.model.images.generation.GenerateImagesRequest;
import com.volcengine.ark.runtime.model.images.generation.ImagesResponse;
import com.volcengine.ark.runtime.service.ArkService;

import cn.hutool.core.lang.UUID;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;

@CrossOrigin(origins = "*")
@RestController
@RequestMapping("/api/doubao/ai")
public class AIDoubaoController {
	private final static Logger LOG = LoggerFactory.getLogger(AIDoubaoController.class);
	private final static String ApiKey = "14706c1d-143b-4e12-8aea-fb5539a441a5";//替换成自己的 API Key
	
	@PostMapping(value = "/ask", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter ask() {
        SseEmitter emitter = new SseEmitter(60_000L); // 超时时间 60 秒

        String content = "你是一个AI表格助手，请根据用户要求生成表格，输出格式为Markdown格式，下面请帮我输出一个车辆维护记录表，请填充好示例数据。";//替换成自己的 API Key
        String model = "doubao-1-5-pro-32k-250115";//替换成自己的 模型id
        // 创建ArkService实例
        ArkService arkService = ArkService.builder().apiKey(ApiKey).build();
        // 模拟流式数据
        new Thread(() -> {
            try {
                // 初始化消息列表
                List<ChatMessage> chatMessages = new ArrayList<>();

                // 创建用户消息
                ChatMessage userMessage = ChatMessage.builder()
                    .role(ChatMessageRole.USER) // 设置消息角色为用户
                    .content(content) // //替换成自己的 提问内容
                    .build();
                
//				final List<ChatMessage> messages = Arrays.asList(
//						ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build(),
//						ChatMessage.builder().role(ChatMessageRole.USER).content("花椰菜是什么？").build(),
//						ChatMessage.builder().role(ChatMessageRole.ASSISTANT).content("花椰菜又称菜花、花菜，是一种常见的蔬菜。").build(),
//						ChatMessage.builder().role(ChatMessageRole.USER).content("再详细点").build());

                // 将用户消息添加到消息列表
                chatMessages.add(userMessage);

                // 创建聊天完成请求
                ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(model)// 需要替换为Model ID
                    .messages(chatMessages) // 设置消息列表
                    .stream(true)//以流式返回
                    .build();

                // 发送聊天完成请求并打印响应
                // 获取响应并打印每个选择的消息内容
                String traceId = UUID.fastUUID().toString(true);
                StringBuilder sb = new StringBuilder();
                arkService.streamChatCompletion(chatCompletionRequest).blockingIterable().forEach(j -> j.getChoices()
                    .forEach(choice -> {
                            //System.out.print(choice.getMessage().getContent());
                            AIChatDTO aiChatDTO = new AIChatDTO();
                            aiChatDTO.setContent(choice.getMessage().getContent().toString());
                            try {
                                emitter.send(SseEmitter.event()
                                    .id(traceId).data(JSON.toJSONString(aiChatDTO)));
                                sb.append(aiChatDTO.getContent());
                            } catch (IOException e) {
                                LOG.error("ask IOException");
                            }
                        }
                    ));
                System.out.println(sb);
                emitter.complete();
            } catch (Exception ex) {
                emitter.completeWithError(ex);
            } finally {
                // 关闭服务执行器
                arkService.shutdownExecutor();
            }
        }).start();

        return emitter;
    }
	
	@PostMapping(value = "/generateImage")
	public ResponseEntity<?> generateImage() {
		String model = "doubao-seedream-3-0-t2i-250415";
		ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
		Dispatcher dispatcher = new Dispatcher();
		ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).apiKey(ApiKey).build();
		try {
			GenerateImagesRequest generateRequest = GenerateImagesRequest.builder()
					.model(model)
					/**
					  1024x1024 （1:1）
					  864x1152 （3:4）
					  1152x864 （4:3）
					  1280x720 （16:9）
					  720x1280 （9:16）
					  832x1248 （2:3）
					  1248x832 （3:2）
					  1512x648 （21:9）
					 */
					.size("1024x1024") // 默认值1024x1024
					.watermark(false)  // 是否在图片右下角添加“AI生成”字样的水印标识
					.prompt("画一个拿金箍棒的哪吒，金箍棒架在肩膀上，嘴上抽着一支中华烟，踩着哥斯拉说：就你还来搞事情？")
					.build();
			ImagesResponse imagesResponse = service.generateImages(generateRequest);
			return new ResponseEntity<>(imagesResponse, HttpStatus.OK);
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			service.shutdownExecutor();
		}
		return new ResponseEntity<>("Empty Response", HttpStatus.OK);
    }
}
